#include "RelocaliserSCARF.h"
#include "Combination.h"
#include "Dbscan.h"
#include "KeyFrame.h"
#include "MapMaker.h"
#include "RandomHelpers.h"
#include <cvd/utility.h>
#include <limits>
#include <sys/time.h>
#include <TooN/SVD.h>
#include <TooN/TooN.h>
#include <math.h>
#include <boost/ptr_container/ptr_vector.hpp>

using namespace CVD;
using namespace std;
using namespace GVars3;

template<typename T> inline bool isinf(T value)
{
  return std::numeric_limits<T>::has_infinity() && value == std::numeric_limits<T>::infinity();
}

double generateSample(double sigma)
{
  double fac, r, v1, v2;
	static bool generate_sample = true;
	static double stored_sample = 0;
	
  if (generate_sample)
  {
    do 
    {
	    v1 = (2.0 * (double)rand()/RAND_MAX) - 1.0;
	    v2 = (2.0 * (double)rand()/RAND_MAX) - 1.0;
	    r = (v1*v1) + (v2*v2);
    } 
    while (r >= 1.0);

    fac = sqrt(-2.0 * log(r) / r);
    stored_sample = v1*fac;
    generate_sample = false;

    return(sigma*v2*fac);
  } 
  else 
  {
    generate_sample = true;
    return(sigma*stored_sample);
  }
}  

uint16_t RelocaliserSCARF::DistanceSAD(uint8_t desc1[64], uint8_t desc2[64])
{
  uint16_t sums[8];
  __m128i sum;

  __m128i diff0 = _mm_sad_epu8(_mm_loadu_si128((__m128i*)&(desc1[0])), _mm_loadu_si128((__m128i*)&(desc2[0])));
  __m128i diff1 = _mm_sad_epu8(_mm_loadu_si128((__m128i*)&(desc1[16])),_mm_loadu_si128((__m128i*)&(desc2[16])));
  __m128i diff2 = _mm_sad_epu8(_mm_loadu_si128((__m128i*)&(desc1[32])),_mm_loadu_si128((__m128i*)&(desc2[32])));
  __m128i diff3 = _mm_sad_epu8(_mm_loadu_si128((__m128i*)&(desc1[48])),_mm_loadu_si128((__m128i*)&(desc2[48])));

  sum = _mm_adds_epu16(diff0, diff1);
  sum = _mm_adds_epu16(sum, diff2);
  sum = _mm_adds_epu16(sum, diff3);

  _mm_storeu_si128((__m128i*)&sums, sum);

  return (sums[0] + sums[4]);
}

RelocaliserSCARF::RelocaliserSCARF(const ATANCamera &camera)
  : mCamera(camera), mPoseEstimator(camera), mnDescriptorDimension(60)
{
  GV3::Register(mgvnMaxImageFeatures, "Relocaliser.MaxImageFeatures", 800, SILENT); // Number of features in the current keyframe to consider.
  GV3::Register(mgvnMaxLocalKeyFrames, "Relocaliser.MaxLocalKeyFrames", 150, SILENT); // Number of nearby keyframes to consider.
  GV3::Register(mgvnMaxLocalMapPoints, "Relocaliser.MaxLocalMapPoints", 3000, SILENT); // Number of nearby keyframes to consider.
  GV3::Register(mgvnRansacMinObservations, "Relocaliser.RansacMinObservations", 30, SILENT);
  GV3::Register(mgvnRansacMaxHypotheses, "Relocaliser.RansacMaxHypotheses", 500, SILENT);
  GV3::Register(mgvnRansacMaxBlocks, "Relocaliser.RansacMaxBlocks", 50, SILENT);
  GV3::Register(mgvnRansacGroupsPerBlock, "Relocaliser.RansacGroupsPerBlock", 10, SILENT);
  GV3::Register(mgvnRansacGroupSize, "Relocaliser.RansacGroupSize", 10, SILENT);
  GV3::Register(mgvdDbscanEps, "Relocaliser.DbscanEps", 0.1, SILENT); // Same as mapper scale (default distance between first two keyframes if scale unknown)
  GV3::Register(mgvnDbscanMinObs, "Relocaliser.DbscanMinObs", 5, SILENT);
  GV3::Register(mgvnStaleMapFrames, "Relocaliser.StaleMapFrames", 180, SILENT);
  GV3::Register(mgvnGlobalSE3, "System.GlobalSE3", 0, SILENT);
  GV3::Register(mgvnGlobalSO3, "System.GlobalSO3", 0, SILENT);
  GV3::Register(mgvdMaxScarfSAD, "Relocaliser.MaxScarfSAD", 200.0, SILENT);
  GV3::Register(mgvdMaxScarfSSD, "Relocaliser.MaxScarfSSD", 1000.0, SILENT);
  mnMaxImageFeatures = *mgvnMaxImageFeatures;
  mnMaxLocalKeyFrames = *mgvnMaxLocalKeyFrames;  // Cache gvar.
  mnMaxLocalMapPoints = *mgvnMaxLocalMapPoints;  // Cache gvar.
  mnRansacMinObservations = *mgvnRansacMinObservations;  // Cache gvar.
  mnRansacMaxHypotheses = *mgvnRansacMaxHypotheses;  // Cache gvar.
  mnRansacMaxBlocks = *mgvnRansacMaxBlocks;  // Cache gvar.
  mnRansacGroupsPerBlock = *mgvnRansacGroupsPerBlock;  // Cache gvar.
  mnRansacGroupSize = *mgvnRansacGroupSize;  // Cache gvar.
  mdDbscanEps = *mgvdDbscanEps;
  mnDbscanMinObs = *mgvnDbscanMinObs;
  mnStaleMapFrames = *mgvnStaleMapFrames;
  mnGlobalSE3 = *mgvnGlobalSE3;
  mnGlobalSO3 = *mgvnGlobalSO3;
  mdMaxScarfSAD = *mgvdMaxScarfSAD;
  mdMaxScarfSSD = *mgvdMaxScarfSSD;
  mFlannIndex = NULL;
  mFlannDescriptors = NULL;
  mFlannWorldPositions = NULL;
  mpOldMapKF.reset();
  mnStaleMapCount = mnStaleMapFrames;
  mbUsingGlobalMap = false;

  ofsLogger.open("RelocaliserSCARF.log", ios_base::trunc);
};

RelocaliserSCARF::~RelocaliserSCARF()
{
  if(ofsLogger.is_open())
    ofsLogger.close();

  if(mFlannIndex != NULL)
  {
    delete mFlannIndex;
    delete mFlannDescriptors;
    delete mFlannWorldPositions;
  }
}

SE3<> RelocaliserSCARF::BestPose()
{
  return mse3Best;
}

bool RelocaliserSCARF::AttemptRecovery(Map *pMap, boost::shared_ptr<KeyFrame> pCurrentKF, boost::shared_ptr<KeyFrame> pMapKF)
{
  uint64_t start_time, end_time;
  start_time = ClockGetTime();

  mvObservations.clear();
  mvRelocaliserPose.clear();
  for(unsigned int i = 0; i < mvClusters.size(); ++i)
    mvClusters[i].clear();
  mvClusters.clear();

  // Nearest position in map has been updated and at least 30 frames since last update...
  if(pMapKF != mpOldMapKF && mnStaleMapCount > 30)
  {
    msLocalMapPoints.clear();

    boost::shared_lock< boost::shared_mutex > lockKeyFrames(pMap->mKeyFramesLock);

    // Add MapPoints from most recent keyframe
    for(std::map<boost::shared_ptr<MapPoint>, Measurement>::iterator it = pMapKF->mMeasurements.begin(); it != pMapKF->mMeasurements.end(); it++)
      if(it->first->nMEstimatorOutlierCount < 10 && !it->first->bBad)
        msLocalMapPoints.insert(it->first);  // Will not accept duplicates

    //Grab the maximum number of keyframes we can handle (sorted by their proximity to most recent keyframe)
    Vector<3> v3MostRecentPos = pMapKF->se3CfromW.inverse().get_translation();
    std::vector<std::pair<double, boost::shared_ptr<KeyFrame> > > vKeyFrames;
    pMap->mrtAllKeyFrames.NearestNeighbours(v3MostRecentPos.get_data_ptr(), pMapKF, mnMaxLocalKeyFrames, vKeyFrames, 2);

    for(unsigned int i = 0; i < vKeyFrames.size(); ++i)
    {
      for(std::map<boost::shared_ptr<MapPoint>, Measurement>::iterator it = vKeyFrames[i].second->mMeasurements.begin(); it != vKeyFrames[i].second->mMeasurements.end(); it++)
      {
        if(it->first->nMEstimatorOutlierCount < 10 && !it->first->bBad)
          msLocalMapPoints.insert(it->first);  // Will not accept duplicates
      }

      if(msLocalMapPoints.size() >= mnMaxLocalMapPoints)
        break;
    }

    lockKeyFrames.unlock();

    mnStaleMapCount = 0;
    mbUsingGlobalMap = false;
    mpOldMapKF = pMapKF;
  }
  else
    mnStaleMapCount++;

  // The map is old and has not been updated in a while so have to surrender and use the entire map...
  if(!mbUsingGlobalMap && mnStaleMapCount > mnStaleMapFrames)
  {
    if(mFlannIndex != NULL)
    {
      delete mFlannIndex;
      delete mFlannDescriptors;
      delete mFlannWorldPositions;
    }
    msLocalMapPoints.clear();

    boost::shared_lock< boost::shared_mutex > lockMapPoints(pMap->mMapPointsLock);

    // Worst case scenario - Search entire map
    unsigned int nMapFeatures = pMap->mvAllMapPoints.size();
    mFlannDescriptors = new cv::Mat(nMapFeatures, mnDescriptorDimension, CV_32F);
    mFlannWorldPositions = new cv::Mat(nMapFeatures, 3, CV_32F);
    for(unsigned int i = 0; i < nMapFeatures; ++i)
    {
      float* Di = mFlannDescriptors->ptr<float>(i);
      for(unsigned int j = 0; j < mnDescriptorDimension; ++j)
        Di[j] = pMap->mvAllMapPoints[i]->pDescriptor[j];

      float* Pi = mFlannWorldPositions->ptr<float>(i);  // Have to copy as mapmaker may modify map in the meantime
      for(unsigned int j = 0; j < 3; ++j)
        Pi[j] = pMap->mvAllMapPoints[i]->v3WorldPos[j];
    }

    lockMapPoints.unlock();

    mFlannIndex = new cv::flann::Index(*mFlannDescriptors, cv::flann::KDTreeIndexParams(8));
    mbUsingGlobalMap = true;
  }

  // Create a single vector of candidates (instead of one for each level)
  std::vector<Candidate> vCandidates;
  vCandidates.insert(vCandidates.end(), pCurrentKF->aLevels[0].vCandidates.begin(), pCurrentKF->aLevels[0].vCandidates.end());
  vCandidates.insert(vCandidates.end(), pCurrentKF->aLevels[1].vCandidates.begin(), pCurrentKF->aLevels[1].vCandidates.end());
  vCandidates.insert(vCandidates.end(), pCurrentKF->aLevels[2].vCandidates.begin(), pCurrentKF->aLevels[2].vCandidates.end());
  vCandidates.insert(vCandidates.end(), pCurrentKF->aLevels[3].vCandidates.begin(), pCurrentKF->aLevels[3].vCandidates.end());

  // Keep only the best
  size_t mnBestFeatures = min((size_t)mnMaxImageFeatures, vCandidates.size());
  int nNumObservations = 0;

  if(mnBestFeatures > 10)
  {
    std::partial_sort(vCandidates.begin(), vCandidates.begin()+mnBestFeatures, vCandidates.end());
    vCandidates.erase(vCandidates.begin()+mnBestFeatures, vCandidates.end());

    mvObservations.resize(vCandidates.size());

    if(mbUsingGlobalMap) // Global matching using randomized KD-trees
    {
      cv::Mat ImageFeatures(vCandidates.size(), mnDescriptorDimension, CV_32F);
      for(unsigned int i = 0; i < vCandidates.size(); ++i)
      {
        float* Ii = ImageFeatures.ptr<float>(i);
        for(unsigned int j = 0; j < mnDescriptorDimension; ++j)
          Ii[j] = (float)vCandidates[i].descriptor[j];
      }

      // Radius constrained search
      cv::Mat MatchIndices(1, 1, CV_32S);
      cv::Mat MatchDists(1, 1, CV_32F);
      for(int i = 0; i < ImageFeatures.rows; ++i)
      {
        int nMatches = mFlannIndex->radiusSearch(ImageFeatures.row(i), MatchIndices, MatchDists, mdMaxScarfSSD, cv::flann::SearchParams(1600));

        if(nMatches > 0)
        {
          int nMatchIndex = MatchIndices.at<int>(0,0);
          mvObservations[nNumObservations].v2ImagePos = vCandidates[i].v2RootPos;
          mvObservations[nNumObservations].v3Ray = unit(unproject(mCamera.UnProject(vCandidates[i].v2RootPos)));
          mvObservations[nNumObservations].v3WorldPos = makeVector(mFlannWorldPositions->at<float>(nMatchIndex,0), mFlannWorldPositions->at<float>(nMatchIndex,1), mFlannWorldPositions->at<float>(nMatchIndex,2));
          nNumObservations++;
        }
      }

      // Nearest-neighbours constrained search
      /*cv::Mat MatchIndices(vCandidates.size(), 1, CV_32S);
      cv::Mat MatchDists(vCandidates.size(), 1, CV_32F);
      mFlannIndex->knnSearch(ImageFeatures, MatchIndices, MatchDists, 1, cv::flann::SearchParams(1600));

      for(unsigned int i = 0; i < vCandidates.size(); i++)
      {
        if(MatchDists.at<float>(i,0) < mdMaxScarfSSD)
        {
          int nMatchIndex = MatchIndices.at<int>(i,0);
          mvObservations[nNumObservations].v2ImagePos = vCandidates[i].v2RootPos;
          mvObservations[nNumObservations].v3Ray = unit(unproject(mCamera.UnProject(vCandidates[i].v2RootPos)));
          mvObservations[nNumObservations].v3WorldPos = makeVector(mFlannWorldPositions->at<float>(nMatchIndex,0), mFlannWorldPositions->at<float>(nMatchIndex,1), mFlannWorldPositions->at<float>(nMatchIndex,2));
          nNumObservations++;
        }
      }*/
    }
    else // Local matching using brute-force SAD
    {
      for(unsigned int i = 0; i < vCandidates.size(); i++)
      {
        bool bMatched = false;
        uint16_t nMinDistance = (uint16_t)mdMaxScarfSAD;

        for(std::set<boost::shared_ptr<MapPoint> >::iterator it = msLocalMapPoints.begin(); it != msLocalMapPoints.end(); ++it)
        {
          uint16_t nDistance = DistanceSAD(&(vCandidates[i].descriptor[0]), (*it)->pDescriptor);
          if(nDistance < nMinDistance)
          {
            bMatched = true;
            nMinDistance = nDistance;
            mvObservations[nNumObservations].v3WorldPos = (*it)->v3WorldPos;
          }
        }

        if(bMatched)
        {
          mvObservations[nNumObservations].v2ImagePos = vCandidates[i].v2RootPos;
          mvObservations[nNumObservations].v3Ray = unit(unproject(mCamera.UnProject(vCandidates[i].v2RootPos)));
          nNumObservations++;
        }
      }
    }

    // Delete any excess observations
    mvObservations.erase(mvObservations.begin()+nNumObservations, mvObservations.end());

    // TODO:: Before doing RANSAC filter out definite outlier observations, ie. observations that do not belong to a cluster, and also create clusters
    // of points that should be considered together ie. a bunch of mappoints in close proximity and then consider each cluster independently.
    //DBSCAN dbscan(mdDbscanEps, mnDbscanMinObs);
    //dbscan.Cluster(mvObservations, mvClusters);
    std::vector<Observation*> vSingleCluster;
    for(unsigned int i = 0; i < mvObservations.size(); ++i)
      vSingleCluster.push_back(&mvObservations[i]);
    mvClusters.push_back(vSingleCluster);

    //SE3GaussianPDF se3PdfWfromC = pCurrentKF->ComputeGlobalWfromC(pMap->GetFfromW()); //TODO: fix and Restore

    for(unsigned int i = 0; i < mvClusters.size(); ++i)
    {
      if(mvClusters[i].size() >= mnRansacMinObservations)
      {
        PoseHypothesis bestPose;
        if(mPoseEstimator.threePointAbsoluteRANSAC(mnRansacMaxHypotheses, mnRansacMaxBlocks, mnRansacGroupsPerBlock, mnRansacGroupSize, mvClusters[i], bestPose, NULL))
          mvRelocaliserPose.push_back(bestPose);
      }
    }
    std::partial_sort(mvRelocaliserPose.begin(), mvRelocaliserPose.begin()+1, mvRelocaliserPose.end());
  }

  end_time = ClockGetTime();

  ofsLogger << mbUsingGlobalMap << " " << mnStaleMapCount << " " << ClockTimeDifference(start_time, end_time)/1000 << " " << mvRelocaliserPose.size() << " " << vCandidates.size() << " " << msLocalMapPoints.size() << " " << pMap->mvAllMapPoints.size() << endl;

  if(mvRelocaliserPose.size() > 0)
  {
    if(mbUsingGlobalMap)
    {
      mbUsingGlobalMap = false;
      delete mFlannIndex;
      mFlannIndex = NULL;
      delete mFlannDescriptors;
      mFlannDescriptors = NULL;
      delete mFlannWorldPositions;
      mFlannWorldPositions = NULL;
    }

    msLocalMapPoints.clear();
    mse3Best = mvRelocaliserPose[0].se3CfromW;
    mpOldMapKF.reset();
    mnStaleMapCount = mnStaleMapFrames;

    return true;
  }
  else
    return false;
};

